#include <stdio.h>
#include <stdlib.h>
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>

static int decode_save_packet(AVCodecContext* pCodecContext, AVPacket* pPacket, AVFrame* pFrame);
static void save_gray_pic(uint8_t* data, int line_len, int xlen, int ylen, char* filename);

int main(int argc, char **argv){
    const char* video_name = argv[1]; //"./small_bunny_1080p_60fps.mp4";
    AVFormatContext* pFormatContext = avformat_alloc_context();
    if(pFormatContext == NULL){
        printf("avformat_alloc_context error\n");
        return -1;
    }
    printf("hello ffmpeg!\n");
    if(avformat_open_input(&pFormatContext, video_name, NULL, NULL)!=0){
        printf("avformat_open_input error\n");
        return -1;
    }

    printf("Video Info -- format: %s, duration: %ld us\n", pFormatContext->iformat->long_name, pFormatContext->duration);
    if(avformat_find_stream_info(pFormatContext, NULL) < 0){
        printf("avformat_find_stream_info error \n");
        return -1;
    }

    printf("\n");
    AVCodecParameters* p_codecpar = NULL;
    AVStream* p_stream = NULL;
    AVCodec* p_avcodec = NULL;
    AVCodecParameters* p_vid_codecpar = NULL;
    AVCodec* p_vid_avcodec = NULL;
    int video_stream_idx = -1;

    for (int i=0; i<pFormatContext->nb_streams; i++){
        p_stream = pFormatContext->streams[i];
        p_codecpar = p_stream->codecpar;
        printf("Time base before coded: %d / %d\n", p_stream->time_base.num, p_stream->time_base.den);
        printf("Real base frame rate of the stream before coded: %d / %d\n", p_stream->r_frame_rate.num, p_stream->r_frame_rate.den);
        printf("\nStart time: %" PRId64, p_stream->start_time);
        printf("\nDuration: %" PRId64, p_stream->duration);
        printf("\nNb_frames: %" PRId64, p_stream->nb_frames);
        printf("\n");

        p_avcodec = avcodec_find_decoder(p_codecpar->codec_id);

        if (p_avcodec == NULL){
            printf("unsupport codec!\n");
            return -1;
        }

        if (p_codecpar->codec_type == AVMEDIA_TYPE_VIDEO){
            if(video_stream_idx == -1){
                video_stream_idx = i;
                p_vid_codecpar = p_codecpar;
                p_vid_avcodec = p_avcodec;
            }
            printf("Video Resolution: %d x %d\n", p_codecpar->width, p_codecpar->height);
        }else if (p_codecpar->codec_type == AVMEDIA_TYPE_AUDIO ){
            printf("Audio -- channels: %d, sample_rate: %d\n", p_codecpar->channels, p_codecpar->sample_rate);
        }

        printf("AVcodec info: %s, %s, %d\n", p_avcodec->name, p_avcodec->long_name, p_avcodec->id);
        printf("\n");
    }
    p_avcodec = NULL;
    p_stream = NULL;
    p_codecpar = NULL;

    AVCodecContext* pCodecContext = avcodec_alloc_context3(p_vid_avcodec);
    if(pCodecContext == NULL){
        printf("avcodec_alloc_context3 error\n");
        return -1;
    }

    if(avcodec_parameters_to_context(pCodecContext, p_vid_codecpar)<0){
        printf("avcodec_parameters_to_context error\n");
        return -1;
    }

    if(avcodec_open2(pCodecContext, p_vid_avcodec, NULL) < 0){
        printf("avcodec_open2 error\n");
        return -1;
    }

    AVPacket* pPacket = av_packet_alloc();
    if(pPacket == NULL){
        printf("av_packet_alloc error\n");
        return -1;
    }

    AVFrame* pFrame = av_frame_alloc();
    if(pFrame == NULL){
        printf("av_frame_alloc error\n");
        return -1;
    }

    int response = 0;
    int how_many_packets = 8;
    while(av_read_frame(pFormatContext, pPacket) >= 0){
        if(pPacket->stream_index == video_stream_idx){
            response = decode_save_packet(pCodecContext, pPacket, pFrame);
            if(response < 0){
                printf("decode_save_packet response < 0\n");
                break;
            }
            if( --how_many_packets <=0 ){
                printf("processeda all packets\n");
                break;
            }
        }
        av_packet_unref(pPacket);
    }
    av_frame_free(&pFrame);
    av_packet_free(&pPacket);
    avcodec_free_context(&pCodecContext);
    return 0;
}

static int decode_save_packet(AVCodecContext* pCodecContext, AVPacket* pPacket, AVFrame* pFrame){
    int response = avcodec_send_packet(pCodecContext, pPacket);
    if(response < 0){
        printf("Error while sending a packet to the decoder %s\n", av_err2str(response));
        return response;
    }
    while(response >= 0){
        response = avcodec_receive_frame(pCodecContext, pFrame);
        if(response == AVERROR(EAGAIN) || response == AVERROR_EOF || response == AVERROR(EINVAL) || response == AVERROR_INPUT_CHANGED){
            break;
        }else if(response < 0){
            printf("Error while receiving  a frame from the decoder %s\n", av_err2str(response));
            return response;
        }else{
            printf("Frame %d (type= %c, size=%d bytes) key_frame %d [PTS %ld] [DTS %d] \n", 
            pCodecContext->frame_number, 
            av_get_picture_type_char(pFrame->pict_type),
            pFrame->pkt_size,
            pFrame->key_frame,
            pFrame->pts,
            pFrame->coded_picture_number
            );
            
            char frame_filename[1024];
            snprintf(frame_filename, sizeof(frame_filename), "%s-%d.pgm", "frame", pCodecContext->frame_number);
            save_gray_pic(pFrame->data[0], pFrame->linesize[0], pFrame->width, pFrame->height, frame_filename);
        }
    }
    return 0;
}

static void save_gray_pic(uint8_t* data, int line_len, int xlen, int ylen, char* filename){
    FILE* f;
    f = fopen(filename, "w");
    fprintf(f, "P5\n%d %d\n%d\n", xlen, ylen, 255);

    for(int i=0; i<ylen; i++){
        fwrite(data + i * line_len, 1, xlen, f);
    }
    fclose(f);
}